package gold.digger;

import gold.utils.InputUtil;

import java.util.*;
import java.util.List;

/**
 * Created by fanzhenyu02 on 2020/6/27.
 * common problem solver template.
 */
public class LC1090 {
    public long startExecuteTime = System.currentTimeMillis();


    class Solution {
        public int largestValsFromLabels(int[] values, int[] labels, int num_wanted, int use_limit) {
            TreeMap<Integer, List<Integer>> map = new TreeMap<>(Collections.reverseOrder());
            for (int i = 0; i < labels.length; i++) {
                int value = values[i];
                int label = labels[i];
                if (map.containsKey(value)) map.get(value).add(label);
                else {
                    List<Integer> labelList = new ArrayList<>();
                    labelList.add(label);
                    map.put(value, labelList);
                }
            }

            int sum = 0;
            Map<Integer, Integer> count = new LinkedHashMap<>();
            for (Integer value : map.keySet()) {
                if (num_wanted <= 0) break;
                List<Integer> labelList = map.get(value);
                for (Integer label : labelList) {
                    if (count.getOrDefault(label, 0) >= use_limit) continue;
                    count.put(label, count.getOrDefault(label, 0) + 1);
                    if (num_wanted-- <= 0) break;
                    sum += value;
                }
            }

            return sum;
        }
    }


    public void run() {
        Solution solution = new Solution();
        int[] values = InputUtil.toIntegerArray("[2,6,1,2,6]");
        int[] labels = InputUtil.toIntegerArray("[2,2,2,1,1]");
        int num_wanted = 1, use_limit = 1;
        System.out.println(solution.largestValsFromLabels(values, labels, num_wanted, use_limit));
    }

    public static void main(String[] args) throws Exception {
        LC1090 an = new LC1090();
        an.run();

        System.out.println("\ncurrent solution total execute time: " + (System.currentTimeMillis() - an.startExecuteTime) + " ms.");
    }
}
